package com.voiceqsologger.websocket;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.voiceqsologger.config.AsrModeProperties;
import com.voiceqsologger.entity.QsoLog;
import com.voiceqsologger.service.StreamingAsrService;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.BinaryWebSocketHandler;

import java.nio.ByteBuffer;

/**
 * WebSocket 处理：接收音频流与控制消息
 * 协议：
 *  - 文本消息(JSON)：{"type":"start","sampleRate":16000,"lang":"ZH"} / {"type":"stop"}
 *  - 二进制消息：原始PCM 16kHz 16bit mono 数据块
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class StreamRecognizeHandler extends BinaryWebSocketHandler {

    private final ObjectMapper objectMapper = new ObjectMapper();
    private final StreamingAsrService streamingAsrService;
    private final QsoRealtimeOrchestrator orchestrator;
    private final AsrModeProperties modeProperties;
@SneakyThrows
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        JsonNode node = objectMapper.readTree(message.getPayload());
        String type = node.path("type").asText();
        if ("start".equalsIgnoreCase(type)) {
            int sampleRate = node.path("sampleRate").asInt(16000);
            String lang = node.path("lang").asText("ZH");
            // mode 字段可选；若未提供，则服务层会使用全局配置 asr.mode
            String mode = node.has("mode") ? node.path("mode").asText(null) : null;
            streamingAsrService.start(session.getId(), sampleRate, lang, mode);
            session.sendMessage(new TextMessage("{\"event\":\"started\"}"));
        } else if ("stop".equalsIgnoreCase(type)) {
            streamingAsrService.stop(session.getId());
            session.sendMessage(new TextMessage("{\"event\":\"stopped\"}"));
        } else if ("newQso".equalsIgnoreCase(type)) {
            orchestrator.forceStartNew(session.getId());
            session.sendMessage(new TextMessage("{\"event\":\"newQso\"}"));
        }
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        ByteBuffer buffer = message.getPayload();
        byte[] chunk = new byte[buffer.remaining()];
        buffer.get(chunk);
        long startAtHolder[] = new long[1];
        streamingAsrService.append(session.getId(), chunk, text -> {
            try {
                session.sendMessage(new TextMessage(objectMapper.createObjectNode()
                        .put("event", "partial")
                        .put("text", text)
                        .toString()));
            } catch (Exception e) {
                log.error("send partial error", e);
            }
        }, finalText -> {
            try {
                QsoLog snapshot = orchestrator.onFinalText(session.getId(), finalText);
                session.sendMessage(new TextMessage(objectMapper.createObjectNode()
                        .put("event", "final")
                        .put("text", finalText)
                        .toString()));
                // 推送结构化快照与当前模式，便于前端即时展示与刷新
                String mode = modeProperties.getMode();
                long llmCost = orchestrator.getLastLlmCostMs(session.getId());
                session.sendMessage(new TextMessage(objectMapper.createObjectNode()
                        .put("event", "qso_snapshot")
                        .put("mode", mode)
                        .put("callsign", nullSafe(snapshot.getCallsign()))
                        .put("signalReport", nullSafe(snapshot.getSignalReport()))
                        .put("rig", nullSafe(snapshot.getRig()))
                        .put("antenna", nullSafe(snapshot.getAntenna()))
                        .put("power", nullSafe(snapshot.getPower()))
                        .put("height", nullSafe(snapshot.getHeight()))
                        .put("location", nullSafe(snapshot.getLocation()))
                        .put("asrCostMs", (startAtHolder[0] > 0 ? (System.currentTimeMillis()-startAtHolder[0]) : 0))
                        .put("llmCostMs", llmCost)
                        .toString()));
            } catch (Exception e) {
                log.error("send final error", e);
            }
        }, () -> {
            // ASR 开始识别
            startAtHolder[0] = System.currentTimeMillis();
            try {
                session.sendMessage(new TextMessage(objectMapper.createObjectNode()
                        .put("event", "asr_progress")
                        .put("status", "start")
                        .toString()));
            } catch (Exception ignore) {}
        });
    }

    private static String nullSafe(String v) { return v == null ? "" : v; }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        streamingAsrService.stop(session.getId());
        super.afterConnectionClosed(session, status);
    }
}


